from pathlib import Path
from typing import override

import pandas as pd

from ...processors import ClusterStage
from ...utils import OutputManager
from ..interface import Marshaler

__all__ = ["KMeansClusterExcelMarshaler"]


class KMeansClusterExcelMarshaler(Marshaler[ClusterStage]):
    _output_path: Path

    def __init__(self, output_manager: OutputManager) -> None:
        self._output_path = output_manager.filepath("kmeans_clusters.xlsx")

    @override
    def marshal(self, stage: ClusterStage) -> None:
        points = stage.points  # shape: (n_samples, n_features)
        labels = stage.labels  # shape: (n_samples,)
        centers = stage.centers  # shape: (n_clusters, n_features)

        df = pd.DataFrame(points)
        df["Cluster"] = labels

        with pd.ExcelWriter(self._output_path, engine="openpyxl") as writer:
            for cluster_id in range(centers.shape[0]):
                cluster_df = df[df["Cluster"] == cluster_id]
                cluster_df.to_excel(writer, sheet_name=f"Cluster {cluster_id}", index=False)
